'''
Dataset Generator
author: RuanQuanyuan@2022/4/21

Data Format
'name of pic' 'label'
instance:
001.jpg 0
002.jpg 1
...
'''

import cv2
import numpy as np
import torch
from torch import Tensor, tensor
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms as tfs


class DatasetGenerator(Dataset):
    def __init__(self, data_path: str, label_path: str, resize=None):
        self.data = []
        self.label = []

        self.data_path = data_path
        self.resize = resize

        with open(label_path, 'r') as f:
            while True:
                data = f.readline()
                if not data:
                    break
                data = data.strip().split(' ')
                self.data.append(data[0])
                self.label.append(int(data[1]))

    def __getitem__(self, item):
        img = cv2.imread(self.data_path + self.data[item])
        img.astype(float)
        if self.resize:
            img = cv2.resize(img, (self.resize, self.resize)) / 255  # 放缩归一化
        else:
            img = img / 255  # 归一化

        img += np.random.normal(loc=0.0, scale=0.001, size=img.shape,)
        img = tensor(img, dtype=torch.float)
        img = img.permute(2, 0, 1)  # 改变颜色矩阵的位置
        img = tfs.RandomHorizontalFlip()(img)
        img = tfs.RandomRotation(10)(img)
        label = tensor(int(self.label[item]))

        return img, label

    def __len__(self):
        return len(self.data)

